
import math

from dexen.system import constants

class Selection(object):
    def __init__(self):
        pass
    
class EnvironmentalSelection(Selection):
    def __init__(self):
        pass

    def setObjectives(self, *attributes):
        self.objective_attributes = attributes

    def setSelectionRate(self, rate):
        self.selection_rate = rate
    
    def dominates(self, ind1, ind2):
        dominated = False
        for objective in self.objective_attributes:
            if objective[1] == constants.MINIMIZE:
                if getattr(ind2, objective[0]) < getattr(ind1, objective[0]):
                    return False
                if getattr(ind1, objective[0]) < getattr(ind2, objective[0]):
                    dominated = True

            if objective[1] == constants.MAXIMIZE:
                if getattr(ind2, objective[0]) > getattr(ind1, objective[0]):
                    return False
                if getattr(ind1, objective[0]) > getattr(ind2, objective[0]):
                    dominated = True
        if dominated:
            return True
        return False
    
    def _updatePareto(self, new_ind, non_dominated_ones):
        for ind in non_dominated_ones[:]:
            if self.dominates(new_ind, ind):
                non_dominated_ones.remove(ind)
            
    def filterDominatedOnes(self, inds):
        assert len(self.objective_attributes) != 0
        non_dominated_ones = []
        non_dominated_ones.append(inds[0])
        
        for ind1 in inds[:]:
            is_dominated = False
            dominated = False
            for ind2 in non_dominated_ones:
                if self.dominates(ind2, ind1):
                    is_dominated = True
                    break
                if self.dominates(ind1, ind2):
                    dominated = True
                    break

            if dominated:
                self._updatePareto(ind1, non_dominated_ones)
                non_dominated_ones.append(ind1)
                
            if not dominated and not is_dominated:
                non_dominated_ones.append(ind1)
            
        
        while len(inds) > 0:
            inds.pop(0)
        
        inds.extend(non_dominated_ones)    
        return non_dominated_ones

    def select(self, inds):
        self.calculate_strength(inds)
        self.calculate_raw_fitness(inds)
        self.calculate_distances(inds)
        self.calculate_density(inds)
        self.calculate_fitness(inds)
                    
        next_archive = [x for x in inds if x.fitness < 1]
        dominated_ones = [x for x in inds if x.fitness >= 1]
        archive_size = len(next_archive)
        N = int(len(inds) * self.selection_rate)
        
        print "N is ", N
            
        if archive_size < N:
            dominated_ones.sort(key = lambda x:(x.fitness))
            next_archive.extend(dominated_ones[:N - archive_size])
            #print "Archive is less than N"
        elif archive_size > N:
            for ind in inds:
                ind.removed = False
            self.archive_truncation(next_archive, N)
            #print "Archive is greater than N"
    
        assert len(next_archive) == N
        self.pop_keys(inds)
        return dominated_ones, next_archive
    
    def calculate_strength(self, inds):
        for ind1 in inds:
            ind1.strength = 0
            for ind2 in inds:
                if self.dominates(ind1, ind2):
                    ind1.strength += 1
    
    def calculate_raw_fitness(self, inds):
        for ind1 in inds:
            ind1.raw_fitness = 0
            for ind2 in inds:
                if self.dominates(ind2, ind1):
                    ind1.raw_fitness += ind2.strength
    
    def calculate_fitness(self, inds):
        for ind in inds:
            ind.fitness = ind.raw_fitness + ind.density 
    
    def get_distance(self, ind1, ind2):
        result = 0.0
        for objective in self.objective_attributes:
            result += (getattr(ind1, objective[0]) - 
                       getattr(ind2, objective[0]))**2 
    
        result = math.sqrt(result)
        return result
        
    def calculate_distances(self, inds):
        for ind1 in inds:
            ind1.distances = []
            for ind2 in inds:
                ind1.distances.append({"ind" : ind2, "distance" : self.get_distance(ind1, ind2)})
            ind1.distances.sort(key = lambda x:(x['distance']))
    
    def calculate_density(self, inds):
        k = int(math.sqrt(len(inds))) 
        for ind in inds:
            ind.density = 1.0 / (ind.distances[k]['distance'] + 1) #????? 1 or 2 
    
    def pop_keys(self, inds):
        for ind in inds:
            delattr(ind, 'strength')
            delattr(ind, 'raw_fitness')
            delattr(ind, 'fitness')
            delattr(ind, 'density')
            delattr(ind, 'distances')
            if hasattr(ind, 'removed'):
                delattr(ind, 'removed')
    
    def get_distance_k(self, ind, k):
        cnt = 0
        pos = 0
        while cnt < k:
            assert pos < len(ind.distances)
            if ind.distances[pos]['ind'].removed == False and \
                            not(ind.distances[pos]['ind'] is ind):
                cnt += 1
            pos += 1
        return ind.distances[pos-1]['distance']
        
                
    def smaller(self, ind1, ind2, N):
        for k in range(N):
            dist1 = self.get_distance_k(ind1, k)
            dist2 = self.get_distance_k(ind2, k)
            if dist1 > dist2:
                return False
            if dist1 < dist2:
                return True
        return True
    
    def remove_individual(self, min_ind, inds):
        for i, ind in enumerate(inds):
            if ind is min_ind:
                inds.pop(i)
                
    
    def archive_truncation(self, inds, N):
        while (len(inds) > N):
            min_ind = inds[0]
            for ind in inds[1:]:
                if self.smaller(ind, min_ind, len(inds)):
                    min_ind = ind
            min_ind.removed = True
            self.remove_individual(min_ind, inds) 
